import numpy as np

from ..robot_state import RobotState, MotionState, LegState


class GaitScheduler:
    def __init__(self, model_nv, t_swing: float, dt: float) -> None:
        super().__init__()
        self.isIni = False
        self.phi = 0.0
        self.t_swing = t_swing
        self.dt = dt
        self.Fz_threshold = 100.0
        self.Fz_L_m = 0.0
        self.Fz_R_m = 0.0

        self.legState = LegState.RST
        self.legStateNext = LegState.LST
        self.motionState = MotionState.STAND

        self.FLest = np.zeros(6)
        self.FRest = np.zeros(6)
        self.torJoint = np.zeros(12)
        self.enableNextStep = False
        self.touchDown = False

        self.fe_r_pos_W = np.zeros(3)
        self.fe_l_pos_W = np.zeros(3)
        self.swingStartPos_W = np.zeros(3)
        self.posHip_W = np.zeros(3)
        self.posST_W = np.zeros(3)
        self.hip_l_pos_W = np.zeros(3)
        self.hip_r_pos_W = np.zeros(3)
        self.dq = np.zeros(model_nv)
        self.stanceStartPos_W = np.zeros(3)
        self.fe_r_rot_W = np.zeros((3, 3))
        self.fe_l_rot_W = np.zeros((3, 3))

        self.dyn_M = np.zeros((model_nv, model_nv))
        self.dyn_Non = np.zeros(model_nv)
        self.J_l = np.zeros((6, model_nv))
        self.J_r = np.zeros((6, model_nv))
        self.dJ_l = np.zeros((6, model_nv))
        self.dJ_r = np.zeros((6, model_nv))
        self.theta0 = 0.0
        self.model_nv = 0

    def read_state(self, robot_state: RobotState):
        self.model_nv = robot_state.model_nv
        self.torJoint[:] = np.zeros(self.model_nv - 6)
        for i in range(self.model_nv - 6):
            self.torJoint[i] = robot_state.motors_tor_cur[i]
        self.dyn_M[:, :] = robot_state.dyn_M[:, :]
        self.dyn_Non[:] = robot_state.dyn_Non[:]
        self.J_l[:, :] = robot_state.J_l[:, :]
        self.dJ_l[:, :] = robot_state.dJ_l[:, :]
        self.J_r[:, :] = robot_state.J_r[:, :]
        self.dJ_r[:, :] = robot_state.dJ_r[:, :]
        self.Fz_L_m = robot_state.fL[2]
        self.Fz_R_m = robot_state.fR[2]
        self.hip_l_pos_W[:] = robot_state.hip_l_pos_W[:]
        self.hip_r_pos_W[:] = robot_state.hip_r_pos_W[:]
        self.fe_l_pos_W[:] = robot_state.fe_l_pos_W[:]
        self.fe_r_pos_W[:] = robot_state.fe_r_pos_W[:]
        self.fe_l_rot_W[:] = robot_state.fe_l_rot_W[:]
        self.fe_r_rot_W[:] = robot_state.fe_r_rot_W[:]
        self.dq[:] = robot_state.dq[:]
        self.motionState = robot_state.motionState

    def write_state(self, robot_state: RobotState):
        robot_state.tSwing = robot_state.tSwing
        robot_state.swingStartPos_W[:] = self.swingStartPos_W
        robot_state.stanceDesPos_W[:] = self.stanceStartPos_W
        robot_state.posHip_W[:] = self.posHip_W
        robot_state.posST_W[:] = self.posST_W
        robot_state.theta0 = self.theta0
        robot_state.legState = self.legState
        robot_state.legStateNext = self.legStateNext
        robot_state.phi = self.phi
        robot_state.FL_est[:] = self.FLest[:]
        robot_state.FR_est[:] = self.FRest[:]
        if self.legState == LegState.LST:
            robot_state.stance_fe_pos_cur_W[:] = self.fe_l_pos_W
            robot_state.stance_fe_rot_cur_W[:] = self.fe_l_rot_W
        else:
            robot_state.stance_fe_pos_cur_W[:] = self.fe_r_pos_W
            robot_state.stance_fe_rot_cur_W[:] = self.fe_r_rot_W
        robot_state.motionState = self.motionState

    def step(self):
        tauAll = np.zeros(self.model_nv)
        tauAll[6: self.model_nv] = self.torJoint
        self.FLest[:] = -np.linalg.pinv(self.J_l @ np.linalg.inv(self.dyn_M) @ self.J_l.T) @ (
                self.J_l @ np.linalg.inv(self.dyn_M) @ (tauAll - self.dyn_Non) + self.dJ_l @ self.dq)
        self.FRest[:] = -np.linalg.pinv(self.J_r @ np.linalg.inv(self.dyn_M) @ self.J_r.T) @ (
                self.J_r @ np.linalg.inv(self.dyn_M) @ (tauAll - self.dyn_Non) + self.dJ_r @ self.dq)

        dPhi = 0.0

        if self.motionState == MotionState.WALK2STAND:
            self.enableNextStep = False
            if self.touchDown:
                self.motionState = MotionState.STAND

        if self.motionState == MotionState.STAND:
            dPhi = 0.0
            self.phi = 0.0
            self.isIni = False
            self.legState = LegState.RST
            self.enableNextStep = False
        elif self.motionState == MotionState.WALK:
            self.enableNextStep = True
            dPhi = 1.0 / self.t_swing * self.dt
        elif self.motionState == MotionState.WALK2STAND:
            dPhi = 1.0 / self.t_swing * self.dt

        self.phi += dPhi
        if self.enableNextStep:
            self.touchDown = False

        if not self.isIni:
            self.isIni = True
            if self.legState == LegState.LST:
                self.swingStartPos_W[:] = self.fe_r_pos_W
                self.stanceStartPos_W[:] = self.fe_l_pos_W
            else:
                self.swingStartPos_W[:] = self.fe_l_pos_W
                self.stanceStartPos_W[:] = self.fe_r_pos_W

        if self.legState == LegState.LST and self.FRest[2] >= self.Fz_threshold and self.phi >= 0.6:
            if self.enableNextStep:
                self.legState = LegState.RST
                self.swingStartPos_W[:] = self.fe_l_pos_W
                self.stanceStartPos_W[:] = self.fe_r_pos_W
                self.phi = 0

        elif self.legState == LegState.RST and self.FLest[2] >= self.Fz_threshold and self.phi >= 0.6:
            if self.enableNextStep:
                self.legState = LegState.LST
                self.swingStartPos_W[:] = self.fe_r_pos_W
                self.stanceStartPos_W[:] = self.fe_l_pos_W
                self.phi = 0

        if not self.enableNextStep:
            if self.legState == LegState.LST and self.FRest[2] >= 200:
                self.touchDown = True
            if self.legState == LegState.RST and self.FLest[2] >= 200:
                self.touchDown = True

        if self.phi >= 1:
            self.phi = 1

        if self.legState == LegState.LST:
            self.posHip_W[:] = self.hip_r_pos_W[:]
            self.posST_W[:] = self.fe_l_pos_W[:]
            self.theta0 = -np.pi * 0.5
            self.legStateNext = LegState.RST
        else:
            self.posHip_W[:] = self.hip_l_pos_W[:]
            self.posST_W[:] = self.fe_r_pos_W[:]
            self.theta0 = np.pi * 0.5
            self.legStateNext = LegState.LST
